import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize

# ---------------------- 1. 数据准备 ----------------------
# 输入数据（问题二中提供的中国平安股票数据）
data = pd.DataFrame({
    '日期': ['2025/08/08', '2025/08/09', '2025/08/10', '2025/08/11', '2025/08/12',
             '2025/08/13', '2025/08/14', '2025/08/15', '2025/08/16', '2025/08/17',
             '2025/08/18', '2025/08/19', '2025/08/20', '2025/08/21', '2025/08/22',
             '2025/08/23', '2025/08/24', '2025/08/25', '2025/08/26', '2025/08/27',
             '2025/08/28', '2025/08/29', '2025/08/30', '2025/08/31', '2025/09/01',
             '2025/09/02', '2025/09/03', '2025/09/04', '2025/09/05'],
    '开盘': [41.42, 99.34, 47.86, 84.11, 64.07, 38.18, 39.73, 37.16, 10.86, -80.13,
             130.92, -16.54, 92.41, -42.89, 105.22, 107.8, 59.42, 115.88, 167.92, 160.52,
             31.25, 73.51, -17.44, 43.34, 57.15, -25.4, 65.09, 35.7, 80.57],
    '最高': [102.31, 99.34, 85.02, 84.11, 64.07, 40.78, 39.73, 37.16, 10.86, 132.85,
             130.92, 93.53, 92.41, 106.82, 109.72, 107.8, 118.99, 172.92, 167.92, 160.52,
             75.62, 73.51, 44.23, 58.13, 57.15, 66.83, 65.09, 82.65, 80.57],
    '最低': [41.42, 47.08, 47.86, 63.09, 37.66, 38.18, 36.22, 10.61, -80.13, -80.13,
             -16.54, -16.54, -42.89, -42.89, 105.22, 57.7, 59.42, 115.88, 157.29, 30.78,
             31.25, -17.44, -17.44, 43.34, -25.4, -25.4, 35.31, 35.7, -75.69],
    '收盘': [99.34, 47.86, 84.11, 64.07, 38.18, 39.73, 37.16, 10.86, -80.13, 130.92,
             -16.54, 92.41, -42.89, 105.22, 107.8, 59.42, 115.88, 167.92, 160.52, 31.25,
             73.51, -17.44, 43.34, 57.15, -25.4, 65.09, 35.7, 80.57, -75.69]
})
data['日期'] = pd.to_datetime(data['日期'])  # 转换日期格式


# ---------------------- 2. 收益率计算函数 ----------------------
def calculate_returns(data, scheme=1):
    """
    计算两种方案下的每日收益率（考虑交易费用）
    :param data: 包含最高价、最低价的DataFrame
    :param scheme: 1或2，代表两种价格方案
    :return: 每日收益率Series
    """
    transaction_fee = 0.0003  # 交易费率0.03%
    returns = []

    for idx in range(len(data)):
        high = data['最高'].iloc[idx]
        low = data['最低'].iloc[idx]

        # 计算买入价和卖出价
        if scheme == 1:
            buy_price = low
            sell_price = high
        else:  # 方案2：四分之一分位点和四分之三分位点
            price_range = high - low
            buy_price = low + 0.25 * price_range
            sell_price = low + 0.75 * price_range

        # 计算简单收益率（扣除双向交易费用）
        if buy_price <= 0:  # 避免价格为负导致计算错误
            return 0
        raw_return = (sell_price - buy_price) / buy_price
        net_return = raw_return - 2 * transaction_fee  # 买入和卖出各收一次费
        returns.append(net_return)

    return pd.Series(returns, index=data['日期'], name=f'方案{scheme}收益率')


# 计算两种方案的收益率
returns1 = calculate_returns(data, scheme=1)
returns2 = calculate_returns(data, scheme=2)

# 合并为DataFrame（若有多支股票，可在此处添加列）
portfolio_returns = pd.DataFrame({
    '中国平安_方案1': returns1,
    '中国平安_方案2': returns2
})


# ---------------------- 3. 组合优化模型 ----------------------
def portfolio_metrics(weights, mean_returns, cov_matrix, rf):
    """计算组合的预期收益、风险和夏普比率"""
    port_return = np.sum(weights * mean_returns)
    port_risk = np.sqrt(np.dot(weights.T, np.dot(cov_matrix, weights)))
    sharpe_ratio = (port_return - rf) / port_risk
    return port_return, port_risk, sharpe_ratio


def objective(weights, mean_returns, cov_matrix, rf):
    """目标函数：最小化负夏普比率（即最大化夏普比率）"""
    _, _, sharpe = portfolio_metrics(weights, mean_returns, cov_matrix, rf)
    return -sharpe


def optimize_portfolio(returns, rf):
    """求解最优组合权重"""
    n_assets = returns.shape[1]
    mean_returns = returns.mean()
    cov_matrix = returns.cov()

    # 约束条件
    constraints = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})  # 权重和为1
    bounds = tuple((0, 1) for _ in range(n_assets))  # 权重在0到1之间（不允许卖空）
    initial_weights = np.array([1 / n_assets for _ in range(n_assets)])  # 初始权重平均分配

    # 求解优化问题
    solution = minimize(
        objective, initial_weights, args=(mean_returns, cov_matrix, rf),
        method='SLSQP', bounds=bounds, constraints=constraints
    )
    return solution.x, mean_returns, cov_matrix


# 设定无风险利率（年化3%，转换为日利率）
rf_daily = 0.03 / 252  # 一年按252个交易日计算

# 分别对两种方案求解最优权重（此处单支股票权重必为1，多支股票时会自动分配）
weights1, mean1, cov1 = optimize_portfolio(portfolio_returns[['中国平安_方案1']], rf_daily)
weights2, mean2, cov2 = optimize_portfolio(portfolio_returns[['中国平安_方案2']], rf_daily)

# 计算组合指标
return1, risk1, sharpe1 = portfolio_metrics(weights1, mean1, cov1, rf_daily)
return2, risk2, sharpe2 = portfolio_metrics(weights2, mean2, cov2, rf_daily)

# ---------------------- 4. 结果输出与可视化 ----------------------
# 打印关键指标
print("=== 组合投资策略结果对比 ===")
print(f"方案1（买入最低价，卖出最高价）：")
print(f"  预期日收益率：{return1:.6f}（年化：{return1 * 252:.2%}）")
print(f"  收益率标准差（风险）：{risk1:.6f}")
print(f"  夏普比率：{sharpe1:.6f}")
print(f"  最优权重：{weights1}")

print("\n方案2（四分之一分位点买入，四分之三分位点卖出）：")
print(f"  预期日收益率：{return2:.6f}（年化：{return2 * 252:.2%}）")
print(f"  收益率标准差（风险）：{risk2:.6f}")
print(f"  夏普比率：{sharpe2:.6f}")
print(f"  最优权重：{weights2}")

# 绘制每日收益率曲线
plt.figure(figsize=(12, 6))
plt.plot(returns1.index, returns1, label='方案1每日收益率', color='blue', alpha=0.7)
plt.plot(returns2.index, returns2, label='方案2每日收益率', color='red', alpha=0.7)
plt.axhline(y=0, color='black', linestyle='--', alpha=0.3)
plt.title('两种方案下的每日收益率对比', fontsize=14)
plt.xlabel('日期', fontsize=12)
plt.ylabel('收益率', fontsize=12)
plt.legend()
plt.grid(alpha=0.3)
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

# 绘制累计收益率曲线
cum_return1 = (1 + returns1).cumprod() - 1
cum_return2 = (1 + returns2).cumprod() - 1

plt.figure(figsize=(12, 6))
plt.plot(cum_return1.index, cum_return1, label='方案1累计收益率', color='blue')
plt.plot(cum_return2.index, cum_return2, label='方案2累计收益率', color='red')
plt.axhline(y=0, color='black', linestyle='--', alpha=0.3)
plt.title('两种方案下的累计收益率对比', fontsize=14)
plt.xlabel('日期', fontsize=12)
plt.ylabel('累计收益率', fontsize=12)
plt.legend()
plt.grid(alpha=0.3)
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()